import torch

class data_prefetcher:
    def __init__(self, loader):
        self.loader = iter(loader)
        self.stream = torch.cuda.Stream()
        self.next_input1 = None
        self.next_input2 = None
        self.next_input3 = None
        self.next_input4 = None
        self.next_target = None
        self.preload()

    def preload(self):
        try:
            self.next_input1, self.next_input2, self.next_input3, self.next_input4, self.next_target = next(self.loader)
        except StopIteration:
            self.next_input1 = None
            self.next_input2 = None
            self.next_input3 = None
            self.next_input4 = None
            self.next_target = None
            return
        with torch.cuda.stream(self.stream):
            self.next_input1 = self.next_input1.cuda(non_blocking=True)
            self.next_input2 = self.next_input2.cuda(non_blocking=True)
            self.next_input3 = self.next_input3.cuda(non_blocking=True)
            self.next_input4 = self.next_input4
            self.next_target = self.next_target.cuda(non_blocking=True)

    def next(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        input1 = self.next_input1
        input2 = self.next_input2
        input3 = self.next_input3
        input4 = self.next_input4
        target = self.next_target
        self.preload()
        return input1, input2, input3, input4, target
    
    def reset(self):
        self.loader = iter(self.loader)
        self.preload()